import os
import oneflow as torch
import argparse
import random
# import torchaudio
import scipy.io.wavfile as siw
import python_speech_features as psf
import numpy as np
from tqdm import tqdm

parser = argparse.ArgumentParser()
parser.add_argument('-f', '--wavscp', type=str, default=None)
parser.add_argument('-n', '--num_mel_bins', type=str, default=40)
parser.add_argument('-p', '--percent_of_total_sents', type=float, default=1.0)
parser.add_argument('-o', '--outputdir', type=str, default=None)
args = parser.parse_args()


global_mean = None
global_var = None
num_sentences = 0

random.seed(1234)

with open(args.wavscp, 'r', encoding='utf-8') as f:
    lines = f.readlines()
    random.shuffle(lines)

    for i in tqdm(range(int(args.percent_of_total_sents * len(lines)))):
        path = lines[i].strip().split()[-1]
        # wavform, sample_rate = torchaudio.load_wav(path)
        sample_rate, wavform = siw.read(path) 

        feature_np = psf.base.logfbank(wavform, samplerate=sample_rate, nfilt=args.num_mel_bins)
        feature = torch.tensor(feature_np,dtype=torch.float32)


        # feature = torchaudio.compliance.kaldi.fbank(
        #     wavform, num_mel_bins=args.num_mel_bins,
        #     sample_frequency=sample_rate, dither=0.0
        #     )

        var, mean = torch.var_mean(feature, dim=0)
        num_sentences += 1

        if global_mean is None:
            global_mean = mean
            global_var = var
        else:
            alpha = (num_sentences - 1) / num_sentences
            beta = 1 / num_sentences
            
            old_mean = global_mean.clone()
            global_mean = alpha * old_mean + beta * mean 
            var1 = global_var + torch.pow(global_mean - old_mean, 2)
            var2 = var + torch.pow(global_mean - mean, 2)
            global_var = alpha * var1 + beta * var

global_std = torch.sqrt(global_var)

assert args.outputdir is not None
np.save(os.path.join(args.outputdir, 'global_cmvn.mean'), global_mean.numpy())
np.save(os.path.join(args.outputdir, 'global_cmvn.std'), global_std.numpy())
print('Global Mean and Std are saved as %s' % args.outputdir)

